package immutablecollections;

import immutablecollections.functions.*;
import immutablecollections.misc.*;
import java.io.*;
import java.util.*;

/**
 * An immutable version of `java.util.Set`.
 * 
 * ## Introduction
 * 
 * <img src="doc-files/classes-focus-on-set.png" alt="classes-focus-on-set" style="width: 40em"/></p>
 *  
 * An `ImSet` is a collection of objects with the condition that, if an object `o` belongs
 * to the set then, for any other member `a` of the set:
 * 
 *     a.equals(o) == false
 * 
 * We therefore assume that the objects in the set have "reasonable" implementations of `hashCode()` and `equals()`
 * 
 * A web search for the phrase **always override hashCode when you override equals** will provide more details of
 * the issues involved.
 * 
 * `ImSets` cannot contain `null`.
 * 
 * ## ImSets are Collections
 * 
 * As with many other collections in this library, this class implements the standard Java interface
 * {@link java.util.Collection} (via the class {@link ImCollection}).
 * See {@link ImCollection} for details of how the standard methods of {@link java.util.Collection} work
 * 
 * ## Add/Remove/Find/Replace
 * 
 * You can add/replace elements to the set using {@link #adding} {@link #add} {@link #replace} and remove them
 * using {@link #removing}.
 * 
 * You can test if an element is in the set using {@link #contains} or {@link #find}.
 * 
 * Methods that find/add/replace elements are **`O`**`(log(n))` where n is the size of the set and log is base 2.
 * 
 * ## Replacing Elements
 * 
 * Sometimes it is convenient to be able to add an element to a set even though that set already contains an
 * element that is "equal to it".
 * 
 * For example {@link ImMap} is implemented as a set of {@link ImMap.Entry} objects where an entry has a key and a value.
 * Two {@link ImMap.Entry} objects are equal
 * iff their keys are equal. So for any pair of entries `e1` and `e2`:
 * 
 *     e1.equals(e2) == e1.key.equals(e2.key)
 *  
 * When we want to replace an entry, we want the new entry with its new value to be inserted in the set, even though
 * there is an existing entry with the same key.
 * 
 * To allow this, there are methods {@link #replace} and {@link #add}. The default behaviour when using
 * {@link #adding}  is to *not* replace
 * 
 * ## Implementation
 * 
 * An ImSet is a sorted set of buckets where 
 * elements whose hash codes are the same are stored in the same bucket.
 * 
 * Buckets are sorted on the hash value of their elements.
 * 
 * To find an element, we first find the bucket with the matching hash value and then look through the bucket to
 * determine if the element is present.
 * 
 * For sets of elements with reasonable hash functions, the average bucket size will be one.
 *
 * @see ImSortedSet
 */
public class ImSet<T> extends ImCollection<T>
{
    /**
     * Enumeration containing values `yes` and `no`
     *
     */
    public enum Replace
    {
        /**
         * Replace the element in the set - even if an *equal* element already exists.
         */
        yes,

        /**
         * Don't replace the element in the set if an *equal* element already exists.
         */
        no
    }

    // All the elements with the same hash code are stored in a bucket
    static class Bucket<B> implements Comparable<Bucket<B>>, Serializable
    {
        private static final Object[] EMPTY_ARRAY = new Object[] {};
        final int hashCode;
        final B[] elements;

        Bucket(final int hashCode, final B[] elements)
        {
            this.hashCode = hashCode;
            this.elements = elements;
        }

        @SuppressWarnings("unchecked")
        public Bucket(final int hashCode)
        {
            this(hashCode, (B[]) EMPTY_ARRAY);
        }

        @SuppressWarnings("unchecked")
        public static <C> Bucket<C> newBucket(final int hashCode, final C element)
        {
            return new Bucket<C>(hashCode, (C[]) new Object[] { element });
        }

        public int compareTo(final Bucket<B> other)
        {
            return hashCode < other.hashCode
                    ? -1
                    : (hashCode > other.hashCode
                            ? 1
                            : 0);
        }

        // If an equal element is already in this bucket then it should be
        // replaced
        Bucket<B> add(final B newElement, Replace replace)
        {
            return newBucketAddingElement(newElement, replace);
        }

        private Bucket<B> newBucketAddingElement(final B newElement, Replace replace)
        {

            int index = indexOf(newElement);

            if (index >= 0)
            {
                // newElement (or something equal to it) was in there already
                if (replace == Replace.no)
                    return this;

                // Copy the array
                B[] newElementsArray = Arrays.copyOf(elements, elements.length);

                // Put `newElement` in at the correct slot
                newElementsArray[index] = newElement;

                return new Bucket<B>(hashCode, newElementsArray);
            }
            else
            {
                // `newElement` was not there already. Create an array with an extra slot and copy the old
                // array into it

                // Copy the array
                B[] newElementsArray = Arrays.copyOf(elements, elements.length + 1);

                // Put `newElement` in the last slot
                newElementsArray[elements.length] = newElement;

                return new Bucket<B>(hashCode, newElementsArray);
            }
        }

        @SuppressWarnings("unchecked")
        // If an equals element is not in this bucket then the bucket should
        // not change
        Bucket<B> remove(final B elementToRemove)
        {

            int index = indexOf(elementToRemove);

            // If `elementToRemove` is not in the array then we are done
            if (index < 0)
                return this;

            // So - elementToRemove is in the array

            if (size() == 1)
                return new Bucket<B>(hashCode, (B[]) EMPTY_ARRAY);

            // Create a new array
            // a b c d
            // If we are removing b then after the copy the new array looks like this:
            // a b c
            // Now copy from c to d from the old array into the new array:
            // a c d
            B[] newArray = Arrays.copyOf(elements, elements.length - 1);
            System.arraycopy(elements, index + 1, newArray, index, newArray.length - index);

            return new Bucket<B>(hashCode, newArray);
        }

        B get(final B element)
        {
            int index = indexOf(element);
            return index >= 0
                    ? elements[index]
                    : null;
        }

        @Override
        public String toString()
        {
            return "( " + hashCode + ": " + elements + ")";
        }

        public int size()
        {
            return elements.length;
        }

        private int indexOf(Object newElement)
        {
            for (int i = 0; i < elements.length; i++)
            {
                if (elements[i].equals(newElement))
                {
                    return i;
                }
            }
            return -1;
        }

        public boolean hasEqualElements(Object other)
        {
            if (!(other instanceof Bucket))
                return false;

            Bucket<?> otherBucket = (Bucket<?>) other;

            if (size() != otherBucket.size())
                return false;

            for (Object element : otherBucket.elements)
            {
                if (indexOf(element) == -1)
                    return false;
            }
            return true;
        }

        @Override
        public int hashCode()
        {
            return hashCode;
        }
    }

    /**
     * The sorted set that contains the buckets
     */
    final ImSortedSet<Bucket<T>> sortedSet;

    /**
     * The number of elements in this set
     */
    final private int size;

    protected String getStats()
    {
        int total = 0;
        int totalElements = 0;
        int max = 0;
        int maxElements = 0;
        for (Bucket<T> b : sortedSet)
        {
            total += b.size();
            totalElements += b.elements.length;
            max = Math.max(max, b.size());
            maxElements = Math.max(maxElements, b.elements.length);
        }

        return //
        "# buckets = " + sortedSet.size() + //
                "\nmax bucket size = " + max + //
                "\ntotal  = " + total + //
                "\nheight  = " + sortedSet.tree.getHeight() + //
                "\naverage bucket size = " + total / (float) sortedSet.size() + //
                "\nmaxElements  = " + maxElements + //
                "\ntotalElements  = " + totalElements + //
                "\naverage elements length = " + totalElements / (float) sortedSet.size();
    }

    @SuppressWarnings({ "unchecked", "rawtypes" })
    private static final ImSet<?> empty = new ImSet(ImSortedSet.empty(), 0);

    /**
     * The (singleton) empty set.
     * 
     */
    @SuppressWarnings("unchecked")
    public static <TT> ImSet<TT> empty()
    {
        return (ImSet<TT>) empty;
    }

    private ImSet(final ImSortedSet<Bucket<T>> node, final int size)
    {
        this.sortedSet = node;
        this.size = size;
    }

    /**
     * Add `elementToAdd` to the set.
     * 
     * The set which contains the same elements as `this` and also contains `elementToAdd`.
     * 
     * If an *equal* element is already in this set then a new set will not be created - instead
     * `this` will be returned.
     */
    public ImSet<T> adding(final T elementToAdd)
    {
        return add(elementToAdd, Replace.no);
    }

    /**
     * Remove `elementToRemove` from the set.
     * 
     * The set which contains the same elements as `this` but does not contain `elementToRemove`.
     * 
     * If `elementToRemove` is not in the set then a new set will not be created - instead
     * `this` will be returned.
     */
    public ImSet<T> removing(final T elementToRemove)
    {
        final Bucket<T> b = getBucketContaining(elementToRemove);

        return b == null
                ? this
                : removeElementFromBucket(elementToRemove, b);
    }

    /**
     * Add `newElement` to the set.
     * 
     * If an *equal* element is already in the set then it will be replaced by `newElement`
     */
    public ImSet<T> replace(final T newElement)
    {
        return add(newElement, Replace.yes);
    }

    /**
     * Add `element` to the set.
     * 
     * @throws NullPointerException if `element == null`
     * 
     */
    public ImSet<T> add(final T newElement, Replace replace)
    {
        NullCheck.check(newElement);

        return addAtPath(newElement, replace, getPathOn(newElement));
    }

    protected ImTreeZipper<Bucket<T>> getPathOn(final T newElement)
    {
        // Try to find the path of a bucket with the correct hash in the tree
        return ImTreeZipper.find(ImTreeZipper.onRoot(sortedSet.tree), bucketWith(newElement));
    }

    protected ImSet<T> addAtPath(final T newElement, Replace replace, final ImTreeZipper<Bucket<T>> path)
    {
        final Bucket<T> bucketFound = path.getFocus().getElement();

        if (bucketFound == null)
        {
            // There wasn't even a bucket with the right hash
            ImTree<Bucket<T>> tree = path.replaceEmptyNode(ImTree.newLeaf(bucketWith(newElement))).close();
            return new ImSet<T>(new ImSortedSet<Bucket<T>>(tree), size + 1);
        }
        else
        {
            // There was a bucket with the right hash - let's try adding `newElement`
            Bucket<T> newBucket = bucketFound.add(newElement, replace);

            if (replace == Replace.yes)
            {
                int newSize = size - bucketFound.size() + newBucket.size();
                ImTree<Bucket<T>> tree = path.replaceElement(newBucket).close();
                return new ImSet<T>(new ImSortedSet<Bucket<T>>(tree), newSize);
            }
            else if (newBucket == bucketFound)
                return this;
            else
            {
                ImTree<Bucket<T>> tree = path.replaceElement(newBucket).close();
                return new ImSet<T>(new ImSortedSet<Bucket<T>>(tree), size + 1);
            }
        }
    }

    private Bucket<T> bucketWith(final T newElement)
    {
        return Bucket.newBucket(hashCodeOf(newElement), newElement);
    }

    private static <C> int hashCodeOf(final C element)
    {
        return element.hashCode();
    }

    protected ImSet<T> removeElementFromBucket(final T element, final Bucket<T> bucketWithMatchingHashCode)
    {
        // Remove the element from the bucket (it might not be in this bucket of course)
        Bucket<T> newBucket = bucketWithMatchingHashCode.remove(element);

        // If the bucket has not changed then `element` was not in it so we are done
        if (newBucket == bucketWithMatchingHashCode)
            return this;
        else
            // If the size of the bucket is now zero then we must remove the
            // bucket
            // otherwise we must add this new bucket
            return new ImSet<T>(newBucket.size() == 0
                    ? sortedSet.remove(bucketWithMatchingHashCode)
                    : sortedSet.adding(newBucket), size - 1);
    }

    protected Bucket<T> getBucketContaining(final T element)
    {
        // Find the bucket with the matching hash code
        return (Bucket<T>) sortedSet.find(bucketWith(element));
    }

    /**
     * If an element that *equals* `elementToFind` is in the set then return it else return `null`.
     * 
     */
    public T find(final T elementToFind)
    {
        final int hashCode = hashCodeOf(elementToFind);

        // Create an entry that we can search for
        final Bucket<T> entry = new Bucket<T>(hashCode);

        // Find it in the tree
        final Bucket<T> foundOrNull = (Bucket<T>) sortedSet.find(entry);

        return foundOrNull == null
                ? null
                : foundOrNull.get(elementToFind);
    }

    /**
     * The set whose elements are obtained from `array`.
     * 
     * ImCollections can't contain `null` so none of the elements in `array` can be `null`
     * 
     * ### Examples:
     * 
     *     on(1, 2, 3, 2)  => [1, 2, 3]
     *     on()            => []
     *     on(1, null)     => throws java.lang.NullPointerException: ImCollections can't contain nulls
     *
     * 
     * @throws NullPointerException if any of the elements in `array` is `null`
     * @see #onAll(Collection)
     * @see #onIterator(Iterator)
     * @see #on
    */
    public static <T> ImSet<T> onArray(final T... array)
    {
        return onIterator(ArrayIterator.on(array));
    }

    /**
     * The set whose elements are obtained from `elementsCollection`.
     * 
     * ImCollections can't contain `null` so none of the elements in `elementsCollection` can be `null`
     * 
     * ### Examples:
     * 
     *     on(Arrays.asList(1, 2, 1))    => [1, 2]
     *     on(Arrays.asList())           => []
     *     on(Arrays.asList(1, null, 3)) => throws java.lang.NullPointerException: ImCollections can't contain nulls
     * 
     * @throws NullPointerException if any of the elements in `elementsCollection` is `null`
     * @see #onIterator(Iterator)
     * @see #onArray(Object...)
     * @see #on
     */
    @SuppressWarnings("unchecked")
    public static <T> ImSet<T> onAll(Collection<? extends T> elementsCollection)
    {
        return elementsCollection instanceof ImSet
                ? (ImSet<T>) elementsCollection
                : onIterator(elementsCollection.iterator());
    }

    /**
     * The set whose elements are obtained by iterating over `iterator`.
     * 
     * ImCollections can't contain `null` so none of the elements in `iterator` can be `null`
     * 
     * ### Examples:
     * 
     *     on(Arrays.asList(1, 2, 3).iterator())     =>  [1, 2, 3]
     *     on(Arrays.asList().iterator())            =>  []
     *     on(Arrays.asList(1, null, 3).iterator())  =>  throws java.lang.NullPointerException: ImCollections can't contain nulls
     * 
     * @throws NullPointerException if any of the elements in `iterator` is `null`
     * @see #onAll(Collection)
     * @see #onArray(Object...)
     * @see #on
     */
    public static <T> ImSet<T> onIterator(Iterator<? extends T> iterator)
    {
        ImSet<T> s = ImSet.empty();

        while (iterator.hasNext())
            s = s.adding(iterator.next());

        return s;
    }

    /**
     * The set whose only element is `element`.
     * 
     * This is convenient when we want to create a list whose only element is an array.
     * 
     * ImCollections can't contain `null` so `element` can't be `null`.
     *     
     * @throws NullPointerException if element is `null`
     * @see #onArray(Object...)
     * @see #onIterator(Iterator)
     * @see #onArray(Object[])
     */
    public static <T> ImSet<T> on(T element)
    {
        return ImSet.<T> empty().adding(element);
    }

    /**
     * The number of elements in `this`.
     */
    public int size()
    {
        return size;
    }

    /**
     * A string representation of `this`.
     */
    @Override
    public String toString()
    {
        return "[" + TextUtils.join(iterator(), ", ") + "]";
    }

    //
    //
    // Paths and Iterators
    //
    //

    static class ImSetIterator<T> implements Iterator<T>
    {
        final private ImTreeIterator<Bucket<T>> oTreeIterator;
        private ArrayIterator<T> elementsIterator;

        public ImSetIterator(final ImSortedSet<ImSet.Bucket<T>> oTree)
        {
            oTreeIterator = oTree.iterator();

            if (oTreeIterator.hasNext())
            {
                final ImSet.Bucket<T> bucket = oTreeIterator.next();
                elementsIterator = new ArrayIterator<T>(bucket.elements);
            }
            else
            {
                elementsIterator = new ArrayIterator<T>();
            }
        }

        public boolean hasNext()
        {
            return elementsIterator.hasNext()
                    ? true
                    : oTreeIterator.hasNext();
        }

        public T next()
        {
            if (!hasNext())
                throw new NoSuchElementException();

            if (elementsIterator.hasNext())
            {
                return elementsIterator.next();
            }

            // There *must* be another bucket because hasNext() is true and buckets can't be empty
            elementsIterator = new ArrayIterator<T>(oTreeIterator.next().elements);

            return elementsIterator.next();
        }

        /**
         * @deprecated Unsupported operation since all `ImCollections` are immutable.
         */
        public void remove()
        {
            throw new UnsupportedOperationException();
        }
    }

    /**
     * An iterator over the elements in `this`.
     *
     * Note that the iterator returned by this method will throw an
     * `UnsupportedOperationException` in response to its
     * `remove` method.
     *
     */
    public Iterator<T> iterator()
    {
        return new ImSetIterator<T>(sortedSet);
    }

    /**
     * The set where each element is the result of evaluating the single argument function `fn` on
     * an element of `this`.
     * 
     * This new set may be smaller than the original set. 
     */
    public <O> ImSet<O> map(Function1<O> fn)
    {
        ImSet<O> result = ImSet.empty();

        for (T thing : this)
        {
            result = result.adding(fn.invoke(thing));
        }

        return result;
    }

    /**
     * `true` if `this` contains `elementToLookFor`.
     * 
     * More formally, returns `true` if and only if `this` contains
     * at least one element `e` such that
     * 
     *     e.equals(elementToLookFor)
     *
     */
    @SuppressWarnings("unchecked")
    @Override
    public boolean contains(Object elementToLookFor)
    {
        return find((T) elementToLookFor) != null;
    }

    /** 
     * Compares `other` with `this` for equality.
     * 
     * Returns `true` if the specified object is also an `ImSet`, both
     * sets have the same size, and:
     * 
     *     this.containsAll(other)
     * 
     * In other words, two sets are defined to be
     * *equal* if they contain *equal* elements.
     * 
     * In *other* other words - this is very similar to the standard Java Set `equals`.
     *
     */
    @SuppressWarnings("rawtypes")
    @Override
    public boolean equals(Object other)
    {
        return this == other
                ? true
                : other instanceof ImSet
                        ? eq((ImSet) other)
                        : false;
    }

    private boolean eq(ImSet<?> otherSet)
    {
        return size() == otherSet.size() && hashCode() == otherSet.hashCode()
                ? elementsEq(sortedSet.iterator(), otherSet.sortedSet.iterator())
                : false;
    }

    private boolean elementsEq(ImTreeIterator<Bucket<T>> itOne, ImTreeIterator<?> itTwo)
    {
        while (itOne.hasNext())
        {
            if (!itOne.next().hasEqualElements(itTwo.next()))
                return false;
        }

        return true;
    }

    @Override
    public int hashCode()
    {
        // TODO - may-2012 - Van - could improve this to account for every object
        // rather than each bucket as it is now
        return sortedSet.hashCode();
    }

    public ImSet<T> addingAll(Collection<? extends T> elementsToAdd)
    {
        ImSet<T> result = this;

        for (T a : elementsToAdd)
        {
            result = result.adding(a);
        }

        return result;
    }

    /**
     * Allows `ImSet<U>` to be cast to `ImSet<T>` (where U is a subtype of T) without generating a compiler warning.
     * 
     * For a detailed description see {@link ImList#upCast()}
     */
    @SuppressWarnings("unchecked")
    public <U> ImSet<U> upCast()
    {
        return (ImSet<U>) this;
    }

    /**
     * The ImSet formed out of the elements of each collection in `collections` in order.
     * 
     * ImCollections can't contain `null` so none of the elements can be `null`
     * 
     * ### Examples:
     * 
     *     List<Number> threeFive = Arrays.<Number> asList(3, 5);
     *     ImSet<Integer> oneTwo = ImSet.on(1, 2);
     * 
     *     joinArray(oneTwo, threeFive)          =>  [1, 2, 3, 5]");
     *     joinArray(oneTwo, threeFive, oneTwo)  =>  [1, 2, 3, 5, 1, 2]
     *     joinArray(on(), on())                 =>  []
     * 
     * @see #addingAll(Collection)
     * @see #joinIterator(Iterator)
     * @see #join(Collection)
     */
    public static <T> ImSet<T> joinArray(Collection<? extends T>... collections)
    {
        return joinIterator(ArrayIterator.on(collections));
    }

    /**
     * The ImSet formed out of the elements of each collection in `iterator` in order.
     * 
     * ImCollections can't contain `null` so none of the elements can be `null`
     * 
     * @see #addingAll(Collection)
     * @see #joinArray(Collection...)
     * @see #join(Collection)
     */
    public static <T> ImSet<T> joinIterator(Iterator<Collection<? extends T>> iterator)
    {
        ImSet<T> concat = ImSet.empty();

        while (iterator.hasNext())
            concat = concat.addingAll(ImSet.onAll(iterator.next()));

        return concat;
    }

    /**
     * The ImSet formed out of the elements of each collection in `collectionOfCollections` in order.
     * 
     * ImCollections can't contain `null` so none of the elements can be `null`
     * 
     * @see #addingAll(Collection)
     * @see #joinArray(Collection...)
     * @see #joinIterator(Iterator)
     */
    public static <T> ImSet<T> join(Collection<? extends Collection<? extends T>> collectionOfCollections)
    {
        return joinIterator((Iterator<Collection<? extends T>>) collectionOfCollections.iterator());
    }

}
